library(librarian)
shelf(Seurat,
      tidyseurat,
      harmony,
      ggpubr,
      SingleR,
      tidySummarizedExperiment,
      magrittr,
      tidyverse)

#' Mark doublets in Seurat by Scrublet
#'
#' @param sobj A Seurat object.
#'
#' Need sctkPythonInstallConda() first, and configured python env in Rstudio, or singlecellTK::sctkSelectConda()
#' A tibble of doublet estimation will be printed out in the function
#'
#' @return A Seurat object with added metadata: 'scrublet_call' and 'scrublet_score'
#'
#' @examples 
#' 
#' marked_sobj <- mark_doublets(sobj)
#' sobj <- subset(marked_sobj, scrublet_call == 'Singlet')
mark_doublets <- function(sobj) {
  singleCellTK::selectSCTKConda()
  sobj <- sobj |>
    as.SingleCellExperiment() |>
    singleCellTK::runScrublet() |>
    colData() |>
    as_tibble(rownames = "rowname") |>
    select(c(rowname, contains("scrublet"))) |>
    column_to_rownames("rowname") |>
    AddMetaData(sobj, metadata = _)

  tidyseurat::count(sobj, scrublet_call) |> print()

  sobj
}

#' Pre-process seurat in one function
#'
#' @param sobj A Seurat object.
#' @param batch Batch variable in the metadat of `sobj`. Default is "orig.ident".
#' @param skip_norm Skip `NormalizeData()` or not. Useful in TPM data. FALSE by default.
#' 
#' A clustered UMAP will be produced in the process.
#' 
#' @return A normalized, scaled, batch-harmonized, clustered Seurat object with PCA and UMAP DimReduc.
#'
#' @examples
#' sobj <- quick_process_seurat(sobj)
quick_process_seurat <- function(sobj,
                                 batch = "orig.ident", leiden = TRUE,
                                 tsne = FALSE, regress.cc = FALSE,
                                 skip_norm = FALSE, pcs = 20, res = .8) {
  if (!skip_norm) {
    sobj <- NormalizeData(sobj)
  }
  
  sobj <- sobj |>
    FindVariableFeatures()
  
  sobj <- if (regress.cc) {
    CellCycleScoring(sobj, s.features = cc.genes.updated.2019$s.genes,
                     g2m.features = cc.genes.updated.2019$g2m.genes) |>
      mutate(cc.diff = S.Score - G2M.Score) |>
      ScaleData(vars.to.regress = 'cc.diff')
  } else {
    ScaleData(sobj)
  }
  
  sobj <- sobj |>
    RunPCA() 
  
  reduc <- 'pca'
  
  if (length(batch) > 1 || length(unique(pull(sobj,batch))) > 1) {
    sobj <- sobj |>
      RunHarmony(batch, lambda = rep(1,length(batch)))
    
    reduc <- 'harmony'
  }
  
  sobj <- sobj |>
    RunUMAP(reduction = reduc, dims = 1:pcs) |>
    FindNeighbors(reduction = reduc, dims = 1:pcs)
    
  sobj <- if (leiden) {
    FindClusters(sobj, resolution = res, algorithm = 4, method = 'igraph')
  } else {
    FindClusters(sobj, resolution = res)
  }

  if (tsne) {
    sobj <- sobj |> RunTSNE(reduction = reduc, dims = 1:pcs)
  }
  
  DimPlot(sobj, label = TRUE, label.box = TRUE, repel = TRUE,
          cols = DiscretePalette(36)) |> print()

  sobj
}

#' Adjust sparse matrix for drop-like seq specific contamination
#'
#' @param toc A sparse matrix of counts (filtered count matrix) from drop-based seq, e.g., 10X.
#' @param tod A sparse matrix of drops (raw count matrix). Will be auto-estimated from `toc` if not provided.
#' @param cluster A named vector, where entries are the cluster IDs and names are cellIDs. If no names are provided, the order is assumed to match the order in `toc`.
#'
#' @return A adjusted sparse matrix
#' @export
#'
#' @examples
#' library(Seurat)
#' sobj <- read_rds("some_seurat.rds")
#' adj_mat <- adjust_soup_solo(GetAssayData(sobj, 'count'), sobj$seurat_clusters)
adjust_soup_solo <- function(toc, tod = NULL, cluster) {
  if (is.null(tod)) {
    soupProf <- data.frame(
    row.names = rownames(toc),
    est = rowSums(toc) / sum(toc),
    counts = rowSums(toc)
  )

  set.seed(42)
  scNoDrops <- soupProf |>
    SoupChannel(toc, calcSoupProfile = FALSE) |>
    setSoupProfile(soupProfile = soupProf) |>
    setClusters(cluster) |>
    autoEstCont() |>
    adjustCounts(roundToInt = TRUE)
  } else {
    set.seed(42)
    scNoDrops <- tod |>
      SoupChannel(toc) |>
      setClusters(cluster) |>
      autoEstCont() |>
      adjustCounts(roundToInt = TRUE)
  }
}

#' Adjust seurat for soup contamination, always assume no raw matrix is provided
#'
#' @param sobjss input seurat object.
#'
#' @return a de-souped seurat object, without metadata
#' @export
#'
#' @examples 
soupx_seurat <- function(sobjss){
  GetAssayData(sobjss, 'count') %>%
    adjust_soup_solo(cluster = sobjss$seurat_clusters) %>%
    CreateSeuratObject(min.cells = 3, min.features = 200, names.field = 2)
}

#' Annotate Seurat by SingleR
#'
#' @param sobj A clustered Seurat object to be annotated.
#' @param ref A ref SingleCellExperiment object, typically from `celldex` package.
#' @param new_label A character. The new annotation of cell type in output.
#' @param fine_label A boolean. If TRUE, will use `ref$label.fine` to run `SingleR()`, otherwise will use `ref$label.main`. FALSE by default.
#'
#' @return A annotated Seurat object.
#'
#' @examples
#' source('path/to/mod_seurat.R')
#' sobj <- quick_process_seurat(sobj) |>
#' mark_cell_type_singler()
mark_cell_type_singler <- function(sobj,
                                   ref,
                                   fine_label = FALSE,
                                   sc_ref = FALSE,
                                   new_label = "singler_label") {
  
    singler_res <- sobj |>
      as.SingleCellExperiment() |>
      SingleR(
        ref = ref,
        clusters = sobj$seurat_clusters,
        labels = ifelse(rep_along(along = ref$label.main, fine_label), ref$label.fine, ref$label.main),
        de.method = ifelse(sc_ref, 'wilcox', 'classic'),
        BPPARAM = BiocParallel::MulticoreParam(progressbar = TRUE, tasks = 0)
      )
  
  plotScoreHeatmap(singler_res,
    show.labels = TRUE,
    show.pruned = TRUE,
    show_colnames = TRUE
  ) |> show()
  
  if (new_label %in% colnames(sobj@meta.data)) {sobj <- sobj |> select(-all_of(new_label))}
  
  message('Using ', new_label, ' as colname of the assigned cell types...')

  singler_label <- tibble(
    singler_label = singler_res$labels,
    seurat_clusters = as.factor(singler_res@rownames)
  ) |>
    set_names(c(new_label, "seurat_clusters")) |>
    left_join(sobj, y = _)
}

#' Add suffix to colnames of the object
#'
#' useful for single cell count matrix to add sample info to cell barcodes. Will connect with '_'. like 'col1' to 'col1_name'.
#' @param x An object with colnames
#' @param name A character, the suffix to add
#'
#' @return `x` with its colnames suffixed with `name`.
#'
#' @examples
#' head(iris)
#' iris_suffix <- add_name_field(iris, 'my_suffix')
#' head(iris_suffix)
add_name_field <- function(x, name) {
  colnames(x) <- colnames(x) |>
    map_chr(\(x)str_c(name, "_", x))

  x
}

#' Title
#'
#' @param to_latent 
#' @param big_class 
#'
#' @return
#' @export
#'
#' @examples
annotate_latents <- function(to_latent, big_class = 'Cell',
                             logfc.thres = .5){
  markers <- to_latent |>
    FindAllMarkers(only.pos = TRUE, logfc.threshold = logfc.thres) |>
    filter(p_val_adj < .05, str_starts(gene, 'LINC|ENSG', negate = TRUE)) |>
    group_by(gene) |>
    slice_max(avg_log2FC) |>
    group_by(cluster) |>
    slice_max(avg_log2FC)
  
  markers |>
    ungroup() |>
    mutate(latent_cluster = str_c(gene, big_class, cluster, sep = '_'),
           seurat_clusters = cluster) |>
    select(latent_cluster, seurat_clusters) |>
    left_join(to_latent, y = _)}

#' Title
#'
#' @param obc 
#'
#' @return
#' @export
#'
#' @examples
infer_cellschat <- function(obc){
  obc |>
    subsetData() |>
    identifyOverExpressedGenes() |>
    identifyOverExpressedInteractions() |>
    computeCommunProb(population.size = TRUE) |>
    aggregateNet() |>
    computeCommunProbPathway() |>
    netAnalysis_computeCentrality()
}


my36colors <-c('#E5D2DD', '#53A85F', '#F1BB72', '#F3B1A0', '#D6E7A3', '#57C3F3', '#476D87',
               '#E95C59', '#E59CC4', '#AB3282', '#23452F', '#BD956A', '#8C549C', '#585658',
               '#9FA3A8', '#E0D4CA', '#5F3D69', '#C5DEBA', '#58A4C3', '#E4C755', '#F7F398',
               '#AA9A59', '#E63863', '#E39A35', '#C1E6F3', '#6778AE', '#91D0BE', '#B53E2B',
               '#712820', '#DCC1DD', '#CCE0F5',  '#CCC9E6', '#625D9E', '#68A180', '#3A6963',
               '#968175')

seurat_markers <- list(
  'B' = c('MS4A1','CD79A','CD19'),
  'PC' = c('JCHAIN'),
  'CD4' = c('IL7R','CCR7','S100A4'),
  'CD8' = 'CD8A',
  'NK' = c('GNLY','NKG7'),
  'Monocyte' = c('CD14','LYZ','FCGR3A','MS4A7'),
  'DC' = c('FCER1A','CST3'),
  'Mast' = c('TPSAB1','CPA3')
)
